"""
author:dlr123
date:2021年10月30日
"""
import torch

def Rastrigin(data:torch.TensorType,**kwargs) :
    A = kwargs.get('A',10)
    return 10*A + torch.sum(data ** 2 - A*torch.cos(2*torch.pi*data),axis=0)


def MishraBird(data,**kwargs) :
    if len(data) > 2 :
        raise AttributeError('MishraBird的输入数据维度只能是2！')
    part1 = torch.sin(data[1])*torch.exp(torch.pow(1 - torch.cos(data[0]),2))
    part2 = torch.cos(data[0])*torch.exp(torch.pow(1 - torch.sin(data[1]),2))
    part3 = torch.pow(data[0] - data[1],2)
    return part1 + part2 + part3


def Goldsteitorchrice(data,**kwargs) :
    if len(data) > 2 :
        raise AttributeError('Goldsteitorchrice的输入数据维度只能是2！')
    part1 = 1 + (data[0] + data[1] + 1) ** 2*(
            19 - 14*data[0] + 3*data[0] ** 2 - 14*data[1] + 6*data[0]*data[1] + 3*data[1] ** 2)
    part2 = 30 + (2*data[0] - 3*data[1]) ** 2*(
            18 - 32*data[0] + 12*data[0] ** 2 + 48*data[1] - 36*data[0]*data[1] + 27*data[1] ** 2)
    return part1*part2


def Rosenbrock(data,**kwargs) :
    data_i = data[:-1]
    data_j = data[1 :]
    fx = (100 * (data_j - data_i ** 2) ** 2 + torch.pow(1 - data_i, 2)).sum(axis=0)
    return fx #取负表示越小越好

def Sphere(data,**kwargs):
    fx = torch.sum(data**2)
    return fx

if __name__ =="__main__":
    fc = Rosenbrock
    data = [1,0,1,1,0,0,1,1,0,1,1,1,1,1,1,1,1,1,1,2]
    data = torch.tensor(data).unsqueeze(1).repeat(1,20)
    pred_y = fc(data)
    print(pred_y)